{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Effect of synchronization noise on the classic ML estimator\n",
    "\n",
    "This notebook reproduces the toy examples of **Figure 1b** which characterizes the effect of synchronization noise on the classic ML estimator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Loads libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# External libs\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "\n",
    "# Use `tick` to simulate synthetic data and for the baseline `Classic MLE` estimator\n",
    "from tick.hawkes.simulation import SimuHawkesExpKernels\n",
    "from tick.hawkes.inference import HawkesADM4\n",
    "\n",
    "%matplotlib inline\n",
    "plt.rcParams['font.size'] = 16\n",
    "\n",
    "# Graph drawing parameters\n",
    "nx_draw_params = {\n",
    "    'pos':{0: (0,0), 1: (1, 0)}, \n",
    "    'node_size': 2e3, 'node_color': 'w', \n",
    "    'labels': {0: '$N_A$', 1: '$N_B$'},\n",
    "    'edgecolors':'k', 'arrowsize': 50, 'width': 2,\n",
    "    'font_size': 20\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize random seed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(12345678)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 1. Define parameters of the simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "n_nodes = 2         # Number of nodes\n",
    "decay = 1.0         # Expoenential decay\n",
    "end_time = 10e3     # Length of the observation window\n",
    "\n",
    "# Background intensity\n",
    "mu = 0.05\n",
    "baseline = mu * np.ones(n_nodes, dtype='float')\n",
    "\n",
    "# Excitation matrix\n",
    "adjacency = np.array([[0.95, 0.0],\n",
    "                      [0.95, 0.0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('True baseline:')\n",
    "print(baseline.round(2))\n",
    "print('True adjacency:')\n",
    "print(adjacency.round(2))\n",
    "\n",
    "# Draw the graph\n",
    "print('\\nEstimated excitation graph:')\n",
    "graph_true = nx.DiGraph(adjacency.T)\n",
    "plt.figure(figsize=(8, 1))\n",
    "nx.draw_networkx(graph_true, **nx_draw_params)\n",
    "plt.axis('off');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 2. Run the simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build an utility function to generate noisy sample from an MHP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data(adjacency, decay, baseline, end_time, z_true):\n",
    "    \"\"\"\n",
    "    Generate a realization of a MHP with synchronization noise\n",
    "    \"\"\"\n",
    "    # Compute the observation window adjusted to the noise (to avoid edge effects)\n",
    "    adjusted_end_time = end_time + z_true.min() - z_true.max()\n",
    "    adjusted_start_time = z_true.max()\n",
    "    # Simulate a realization of the process\n",
    "    simu = SimuHawkesExpKernels(adjacency=adjacency,\n",
    "                                decays=decay,\n",
    "                                baseline=baseline,\n",
    "                                end_time=adjusted_end_time,\n",
    "                                verbose=False)\n",
    "    simu.simulate()\n",
    "    # Add synchronization noise to it\n",
    "    events = []\n",
    "    for m, orig_events_m in enumerate(simu.timestamps):\n",
    "        events_m = orig_events_m + z_true[m] - adjusted_start_time\n",
    "        events_m = events_m[(events_m >= 0) & (events_m < end_time)]\n",
    "        events.append(events_m)\n",
    "    return events"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the experiments as follows:\n",
    "- First fix $z^A$ to an arbitrary value, e.g. $z^A = 0$.\n",
    "- Then iterate over a range of `N` values of $z^B$ between $-10$ and $10$.\n",
    "- Average the results of each value of $[z^A, z^B]$ over `K` simulations.\n",
    "\n",
    "**WARNING:** This cell is quite computationally expensive and may take some time to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "N = 101  # Number of noise values to choose for `z^B`\n",
    "K = 10   # Number of simulations per value of `z^B`\n",
    "\n",
    "# Set the synchronization noise values in each dimension\n",
    "z_1 = 0.0 # Arbitrarily fix z_1 to zero\n",
    "z_2_range = np.linspace(-10.0, 10.0, N) # Vary `z_2` from -10 to 10\n",
    "\n",
    "# Init the results of the experiments in this array\n",
    "adjacency_arr = np.zeros((N, K, 2, 2))\n",
    "\n",
    "# Iterate of values of `z^B`\n",
    "for i, z_2 in enumerate(z_2_range):\n",
    "    # Set the value of the noise\n",
    "    z = np.hstack((z_1, z_2))\n",
    "    # Run `K` experiments for this value of `z`\n",
    "    for k in range(K):\n",
    "        # Generate a noisy realization of the MHP\n",
    "        events = generate_data(adjacency=adjacency, decay=decay, baseline=baseline, end_time=end_time, z_true=z)\n",
    "        # Apply the classic MLE\n",
    "        learner = HawkesADM4(decay, C=1e3, lasso_nuclear_ratio=1.0, n_threads=2)\n",
    "        learner.fit(events, end_time)\n",
    "        adjacency_arr[i,k] = learner.adjacency\n",
    "        print(f'{i*K + k + 1:d}/{N*K:d} experiments performed...', end='\\r', flush=True)\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 3. Reproduce Figure 1b from the paper\n",
    "\n",
    "Plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 4))\n",
    "plt.grid()\n",
    "plt.errorbar(z_2_range, np.mean(adjacency_arr[:,:,0,1], axis=1), np.std(adjacency_arr[:,:,0,1], axis=1), \n",
    "             c='C0', label=r'Estimated $\\hat{\\alpha}_{AB}(z^B - z^A)$')\n",
    "plt.errorbar(z_2_range, np.mean(adjacency_arr[:,:,1,0], axis=1), np.std(adjacency_arr[:,:,1,0], axis=1), \n",
    "             c='C1', label=r'Estimated $\\hat{\\alpha}_{BA}(z^B - z^A)$')\n",
    "plt.axhline(y=adjacency[0,1], c='C0', ls='--', label=r'Ground truth $\\alpha_{AB}$')\n",
    "plt.axhline(y=adjacency[1,0], c='C1', ls='--', label=r'Ground truth $\\alpha_{BA}$')\n",
    "plt.xlabel(r'$z^B - z^A$')\n",
    "plt.ylabel('Kernel coefficients')\n",
    "plt.legend()\n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
